'''
Created on Jun 4, 2010

@author: oabalbin
'''

import numpy as np
import scipy as sp
from scipy import stats
import matplotlib.pyplot as plt
import matplotlib.mlab as mlab

from collections import deque, defaultdict


class gene_copa:
    def __init__(self,indicator,genename,mean_rank,mean_copa_score,mean_copa_ind,samples,outlier_quality,genelocation):
        self.indicator=indicator
        self.name=genename
        self.location=genelocation
        self.mean_copa_rank=mean_rank
        self.mean_copa_score=mean_copa_score
        self.mean_copa_ind = mean_copa_ind
        self.samples = samples
        self.quality = outlier_quality

class myArray:
    """
    It creates a minimal array with gene expression - samples
    """
    def __init__(self, geneExpressionArray, gene_dictionary, sample_dictionary, normal_samples, gene_interval):
        
        self.expVal = geneExpressionArray   # It is a masked array
        self.gendict = gene_dictionary
        self.sampledict = sample_dictionary
        self.expVal_copainsample=np.empty((geneExpressionArray.shape[0],geneExpressionArray.shape[1]))
        self.expVal_copacohort=np.empty((geneExpressionArray.shape[0],geneExpressionArray.shape[1]))
        self.expVal_foldchange=np.empty((geneExpressionArray.shape[0],geneExpressionArray.shape[1]))
        self.expVal_foldchange_mean=np.empty((geneExpressionArray.shape[0],geneExpressionArray.shape[1]))
        self.expVal_nsaf=np.empty((geneExpressionArray.shape[0],geneExpressionArray.shape[1]))
        
        
        #self.expVal_isoaverage=np.empty((len(set(gene_dictionary.values())),geneExpressionArray.shape[1]))
        
        self.normal_samples = normal_samples
        self.genInterval = gene_interval
    
    def gene_hugo_names(self,tga):
        """
        It creates a dictionary with hugo names translates the refseq names to hugo names
        """
        hugo_names = defaultdict()
        translation_key = defaultdict()
        for genename in self.gendict.values():
            if genename.startswith('NM'):
                hname = tga.translate_refflat2hugo([genename])
                if hname:
                    hname = hname[0]
                    translation_key[genename] = hname
                else:
                    hname = genename
                hugo_names[hname] = 0
            else:
                translation_key[genename] = genename
                hugo_names[genename] = 0
                
        return translation_key, hugo_names
    
    def get_list_of_samples(self):
        print_samples=['z' for i in range(len(self.sampledict))]
        for j, sp in self.sampledict.iteritems():
            #print sp, j
            print_samples[j] = sp
        
        return list(print_samples)
    
    def get_list_of_genes(self):
        gene_names_list= ['z' for i in range(len(self.gendict))]
        for gen,i in self.gendict.iteritems():
            gene_names_list[i]=gen
        
        return list(gene_names_list)

    
    def log2_norm(self, value_shift):
        """
        Log2 transforms the data
        Note: pay attention of what to do when expression value eual 0, increase by 0.1 or 2.
        """
        self.expVal = np.log2(self.expVal + value_shift)
    
    def center_scale_norm(self):
        """
        Normalization of the data.
        It follows oncomine normalization method.
        All data sets are log-transformed and median-centered per array, 
        and standard deviations were normalized to one per array.
        """
        
        #self.expVal = np.log2(self.expVal)
        # Get the mean of each array-sample. So, mean by column 
        exp_mean = np.ma.mean(self.expVal, axis=0)
        # Get the standard deviation of each array-sample. So, mean by column
        exp_std = np.ma.std(self.expVal, axis=0)            
        exp_norm = (self.expVal - exp_mean)/exp_std

        # Normalized data.
        self.expVal = exp_norm
    
    def copa_norm_in_sample(self):
        """
        Normalization of the data.
        It follows oncomine normalization method.
        All data sets are log-transformed and median-centered per array, 
        and standard deviations were normalized to one per array.
        """
        
        #self.expVal = np.log2(self.expVal)
        # In order to normalize by the median population
        gene_median = np.ma.median(self.expVal,axis=1)
        median_col = np.ma.reshape(gene_median,(self.expVal.shape[0],-1))
        gene_median_matrix = np.tile(median_col,self.expVal.shape[1])

        thisexpVal = np.ma.copy(self.expVal) 
        
        thisexpVal=np.ma.true_divide(thisexpVal , gene_median_matrix)
        # Get the mean of each array-sample. So, mean by column 
        sample_median = np.ma.median(self.expVal, axis=0)
        # Get the standard deviation of each array-sample. So, mean by column
        exp_std = np.ma.std(self.expVal, axis=0)   
        # tile for the sample matrix
        #sample_median_matrix = np.tile(sample_median,(self.expVal.shape[0]),1)
        # Get the mad for each sample
        sample_mad = np.median(np.abs(thisexpVal - sample_median))
        
        exp_norm = (self.expVal - sample_median)/sample_mad

        # Normalized data.
        self.expVal_copainsample = exp_norm
    
    def copa_norm_across_samples(self, expression_cutoff):
        """
        Copa analisis accros the different samples
        """
        gene_median = np.ma.median(self.expVal,axis=1)
        median_col = np.ma.reshape(gene_median,(self.expVal.shape[0],-1))
        gene_median_matrix = np.tile(median_col,self.expVal.shape[1])
        # Get the gene mad
        gene_mad = np.ma.median(np.ma.abs(self.expVal - gene_median_matrix),axis=1) + expression_cutoff
        ############# Important ################
        # when gene_mad = 0.0 set it to den_mad = 1.0
        #gene_mad[gene_mad == 0.0] = 1.0
         
        # make a matrix with the mad
        gene_mad_col = np.ma.reshape(gene_mad,(self.expVal.shape[0],-1))
        gene_mad_matrix = np.tile(gene_mad_col,self.expVal.shape[1])
        # Get the copa index for each gene in each cell line
        gene_copa = np.ma.true_divide((self.expVal - gene_median_matrix), gene_mad_matrix)
        
        self.expVal_copacohort = gene_copa
        
        
    def get_the_outliers(self,percentile, gene_outliersample_dict, copa_rank_dict, copa_score_dict):
        """
        Returns the dictionary for each gene and the sample in which it is an outlier
        """
        
        #gene_outliersample_dict=defaultdict()
        #copa_rank_dict = defaultdict()
        indicator = np.array(range(self.expVal_copacohort.shape[0]))
        sample_indicator = np.array(range(self.expVal_copacohort.shape[1]))
        ### Get the percentiles for all genes at once.
        copa_expVal_transpose  =  np.transpose(self.expVal_copacohort)
        copa_score = np.array(sp.stats.scoreatpercentile(copa_expVal_transpose,percentile))
        ### Get the rank for all genes
        copa_rank = np.argsort(copa_score)
        ### Copa dictionary
        nelements=len(indicator)
        for i,k in zip(indicator,copa_rank):
            copa_score_dict[i].append(copa_score[i])
            copa_rank_dict[i].append(k)

        for i in indicator:
            #outlier_samples = indicator[self.expVal_copacohort[i,:] > copa_score[i]]
            outlier_samples = sample_indicator[self.expVal_copacohort[i,:] > copa_score[i]]
            gene_outliersample_dict[i].append(outlier_samples)
        
        return gene_outliersample_dict, copa_rank_dict, copa_score_dict
    
    
    def fold_change_accross_samples_median(self):
        """
        It calculates a fold change value for each gene.
        using as normalizing factor the median across all 
        samples
        """
        gene_median = np.ma.median(self.expVal,axis=1)
        median_col = np.ma.reshape(gene_median,(self.expVal.shape[0],-1))
        gene_median_matrix = np.tile(median_col,self.expVal.shape[1])
        # Get the gene mad
        self.expVal_foldchange = np.ma.true_divide(self.expVal, gene_median_matrix)
        
    def fold_change_accross_samples_mean(self):
        """
        It calculates a fold change value for each gene.
        using as normalizing factor the mean across all 
        samples
        """
        gene_mean = np.ma.mean(self.expVal,axis=1)
        # the minimum mean should be 1. 
        gene_mean[gene_mean<1.0] = 1.0
        mean_col = np.ma.reshape(gene_mean,(self.expVal.shape[0],-1))
        gene_median_matrix = np.tile(mean_col,self.expVal.shape[1])
        # Get the gene mad
        self.expVal_foldchange_mean = np.round(np.ma.true_divide(self.expVal, gene_median_matrix))
        
    def fold_change_accross_samples_nsaf(self):
        """
        It calculates a fold change value for each gene.
        using as normalizing factor the mean across all 
        samples
        """
        gene_mean = np.ma.mean(self.expVal_nsaf,axis=1)
        # the minimum mean should be 1. 
        gene_mean[gene_mean<1.0] = 1.0
        mean_col = np.ma.reshape(gene_mean,(self.expVal_nsaf.shape[0],-1))
        gene_median_matrix = np.tile(mean_col,self.expVal_nsaf.shape[1])
        # Get the gene mad
        self.expVal_foldchange_mean = np.round(np.ma.true_divide(self.expVal_nsaf, gene_median_matrix))
        
    
    
    def get_gene_percentiles(self, percentile):
        """
        It returns an array with the expression value corresponding to ith percentile for gene i 
        """
        gene_percentiles = [sp.stats.scoreatpercentile(self.expVal[i,:],percentile) for i in range(self.expVal.shape[0])]

        return np.array(gene_percentiles)
    
    def sdv_normals(self,gene):
        """
        Calculate the standard deviation for the normal samples
        """
        normals_std = np.std(self.expVal[gene,np.array(self.normal_samples)])
        
        return normals_std
    
    def sdv_tumors_mets(self,gene):
        """
        """
        normal_samples = set(self.normal_samples)
        all_samples = set(self.sampledict.keys())
        tumors_samples = np.array(list(all_samples.difference(normal_samples)))
                
        tumors_std = np.std(self.expVal[gene,tumors_samples])
        
        return tumors_std
    
    def sdv_array(self,expmat):
        """
        calculate the standar deviation for all the array at the gene level
        """
        return np.std(expmat,axis=1)
        
    def median_array(self,expmat):
        """
        Returns the distribution for the median gene expression
        """
        return np.median(expmat,axis=1)
    
    def average_isoforms(self,gene_dictionary):
        """
        It produces a matrix of averaged isoform.
        gene_dictionary : key=gene name, value=list of row with the same gene
        """
        self.expVal_isoaverage=np.empty((len(gene_dictionary),self.expVal.shape[1]))
        gene_list = deque()
        i=0
        for gene, rows in gene_dictionary.iteritems():
            #print len(gene_dictionary), len(set(gene_dictionary.keys())), self.expVal.shape[0], self.expVal_isoaverage.shape[0], i, gene, rows 
            gene_list.append(gene)
            self.expVal_isoaverage[i,:] = np.mean(self.expVal[rows,],axis=0)
            i+=1
        
        return gene_list
      
    def sdv_array_notisoforms(self):
        """
        calculate the standar deviation for all the array at the gene level
        """
        return np.std(self.expVal_isoaverage,axis=1)
        
    def median_array_notisoforms(self):
        """
        Returns the distribution for the median gene expression
        """
        return np.median(self.expVal_isoaverage,axis=1)
    
    
    def nsaf_normalization(self):
        """
        calculates a nsaf normalization for the proteomics arryas
        """
        # normalize over protein length
        prot_lenght = np.array(map(int,self.genInterval))
        length_col = np.ma.reshape(prot_lenght,(self.expVal.shape[0],-1))
        prot_length_matrix = np.tile(length_col,self.expVal.shape[1])
        self.expVal_nsaf = np.ma.true_divide(self.expVal, prot_length_matrix)
        # Sum over proteins in the sample
        sample_sum = np.sum(self.expVal_nsaf,axis=0)
        # Get the final normalization 
        self.expVal_nsaf = np.ma.true_divide(self.expVal_nsaf, sample_sum)
        
    def log2_norm_nsaf(self, value_shift):
        """
        Log2 transforms the data
        Note: pay attention of what to do when expression value eual 0, increase by 0.1 or 2.
        """
        self.expVal_nsaf = np.log2(self.expVal_nsaf + value_shift)
        print self.expVal_nsaf
    
    def normalize_exp_bynormals(self,normal_samples):
        """
        It normalizes the expression value of the array by the median expression value of the normals.
        The modified expVal matrix has the same size than the original one.
        """

        median_expnormals=np.ma.median(self.expVal[:,normal_samples],axis=1)
        median_col = np.reshape(median_expnormals,(self.expVal.shape[0],-1))
        med_new_arr = np.tile(median_col,self.expVal.shape[1])
        
        # Difference between median normal and tumors
        self.expVal = self.expVal - med_new_arr
    
    def convert2count(self, expmat, featlength, libsize):
        '''
        It converts a matrix of RPKM values to count data.
        input: matrix, feature length, library size 
        RPKM = 10^9 x C / NL, which is really just simply C/N
        C= the number of mappable reads that felt onto the gene's exons
        N= total number of mappable reads in the experiment
        L= the sum of the exons in base pairs.
        C = RPKM*N*L/10^9
        '''
        libsizem = np.tile(libsize,(expmat.shape[0]))
        libsizemat = np.reshape(libsizem,(expmat.shape[0],-1))
        featlengthv = np.reshape(featlength,(expmat.shape[0],-1))
        featlengthmat = np.tile(featlengthv,(expmat.shape[1]))
        
        countsmat = np.ceil((expmat*libsizemat*featlengthmat)/float(1e9))
        
        return countsmat
    
    def get_feature_length(self):
        '''
        return the feature length
        '''
        featlength = np.zeros(len(self.genInterval))
        for i,intv in self.genInterval.iteritems():
            chr, loc = intv.split(':')
            start, end = loc.split('-')
            featlength[i] = int(end) - int(start)
        
        return featlength
        
        
        
            
            
        
        



        
 
  
  
    
        
        





































            
    

    

